#define K_MAX_SHAPE_DIM 0
#include "kernel_operator.h"
#include <type_traits>
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;
template<typename TYPE_Y> class KernelSelectV2 {
    using T = TYPE_Y;
public:
    __aicore__ inline KernelSelectV2() {}
    __aicore__ inline void Init(GM_ADDR condition, GM_ADDR x1, GM_ADDR x2, GM_ADDR y,
        uint16_t finalTileNum, uint16_t tileDataNum, uint16_t TailDataNum, TPipe* pipe) {
        this->tileNum = finalTileNum;
        this->tileDataNum = tileDataNum;
        this->tailDataNum = TailDataNum;
        
        conditionGm.SetGlobalBuffer((__gm__ DTYPE_CONDITION*)condition);
        x1Gm.SetGlobalBuffer((__gm__ DTYPE_X1*)x1);
        x2Gm.SetGlobalBuffer((__gm__ DTYPE_X2*)x2);
        yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y);

        pipe->InitBuffer(inQueueCondition, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_CONDITION));
        pipe->InitBuffer(inQueueX1, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_X1));
        pipe->InitBuffer(inQueueX2, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_X2));
        pipe->InitBuffer(outQueueY, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_Y));

        pipe->InitBuffer(QueueTmpbit, this->tileDataNum * sizeof(uint8_t));
        if constexpr (std::is_same_v<T, half> || std::is_same_v<T, int32_t> || std::is_same_v<T, int8_t>) {
            pipe->InitBuffer(QueueTmp1, this->tileDataNum * sizeof(half));
        }
        if constexpr (std::is_same_v<T, int8_t>) {
            pipe->InitBuffer(QueueTmp2, this->tileDataNum * sizeof(half));
        }

    }
    __aicore__ inline void Process() {
        int32_t loopCount = this->tileNum -1;

        this->processDataNum = this->tileDataNum;
        for (int32_t i = 0; i < loopCount; i++) {
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
        this->processDataNum = this->tailDataNum;
        CopyIn(loopCount);
        Compute(loopCount);
        CopyOut(loopCount);
    }
    private:
    __aicore__ inline void CopyIn(int32_t progress) {
        LocalTensor<DTYPE_CONDITION> conditionLocal = inQueueCondition.AllocTensor<DTYPE_CONDITION>();
        DataCopy(conditionLocal, conditionGm[progress * this->tileDataNum], this->processDataNum);
        inQueueCondition.EnQue(conditionLocal);

        LocalTensor<DTYPE_X1> x1Local = inQueueX1.AllocTensor<DTYPE_X1>();
        LocalTensor<DTYPE_X2> x2Local = inQueueX2.AllocTensor<DTYPE_X2>();
        DataCopy(x1Local, x1Gm[progress * this->tileDataNum], this->processDataNum);
        DataCopy(x2Local, x2Gm[progress * this->tileDataNum], this->processDataNum);
        inQueueX1.EnQue(x1Local);
        inQueueX2.EnQue(x2Local);
    }
    __aicore__ inline void Compute(int32_t progress) {
        LocalTensor<uint8_t> conditionLocal = inQueueCondition.DeQue<DTYPE_CONDITION>().ReinterpretCast<uint8_t>();
        LocalTensor<DTYPE_X1> x1Local;
        LocalTensor<DTYPE_X2> x2Local;
        LocalTensor<DTYPE_Y> yLocal = outQueueY.AllocTensor<DTYPE_Y>();
        auto tmpbit = QueueTmpbit.Get<uint8_t>();
        
        if constexpr (std::is_same_v<T, int32_t>) {
            auto tmp1 = QueueTmp1.Get<half>();
            Cast(tmp1, conditionLocal, RoundMode::CAST_NONE, this->processDataNum);
            inQueueCondition.FreeTensor(conditionLocal);

            CompareScalar(tmpbit, tmp1, (half)0, CMPMODE::EQ, this->processDataNum);
            x1Local = inQueueX1.DeQue<DTYPE_X1>();
            x2Local = inQueueX2.DeQue<DTYPE_X2>();
            Select(yLocal.ReinterpretCast<float>(), tmpbit, x2Local.ReinterpretCast<float>(), x1Local.ReinterpretCast<float>(), SELMODE::VSEL_TENSOR_TENSOR_MODE, this->processDataNum);
        }
        else if constexpr (std::is_same_v<T, int8_t>) {
            auto tmp1 = QueueTmp1.Get<half>();
            auto tmp2 = QueueTmp2.Get<half>();

            Cast(tmp1, conditionLocal, RoundMode::CAST_NONE, this->processDataNum);
            inQueueCondition.FreeTensor(conditionLocal);

            CompareScalar(tmpbit, tmp1, (half)0, CMPMODE::EQ, this->processDataNum);
            x1Local = inQueueX1.DeQue<DTYPE_X1>();
            x2Local = inQueueX2.DeQue<DTYPE_X2>();
            Cast(tmp1, x1Local, RoundMode::CAST_NONE, this->processDataNum);
            Cast(tmp2, x2Local, RoundMode::CAST_NONE, this->processDataNum);
            Select(tmp1, tmpbit, tmp2, tmp1, SELMODE::VSEL_TENSOR_TENSOR_MODE, this->processDataNum);
            Cast(yLocal, tmp1, RoundMode::CAST_NONE, this->processDataNum);
        }
        else if constexpr (std::is_same_v<T, half>) {
            auto tmp1 = QueueTmp1.Get<half>();
            Cast(tmp1, conditionLocal, RoundMode::CAST_NONE, this->processDataNum);
            inQueueCondition.FreeTensor(conditionLocal);

            CompareScalar(tmpbit, tmp1, (half)0, CMPMODE::EQ, this->processDataNum);
            x1Local = inQueueX1.DeQue<DTYPE_X1>();
            x2Local = inQueueX2.DeQue<DTYPE_X2>();
            Select(yLocal, tmpbit, x2Local, x1Local, SELMODE::VSEL_TENSOR_TENSOR_MODE, this->processDataNum);
        }
        else //float
        {
            Cast(yLocal.ReinterpretCast<half>(), conditionLocal, RoundMode::CAST_NONE, this->processDataNum);
            inQueueCondition.FreeTensor(conditionLocal);

            CompareScalar(tmpbit, yLocal.ReinterpretCast<half>(), (half)0, CMPMODE::EQ, this->processDataNum);
            x1Local = inQueueX1.DeQue<DTYPE_X1>();
            x2Local = inQueueX2.DeQue<DTYPE_X2>();
            Select(yLocal, tmpbit, x2Local, x1Local, SELMODE::VSEL_TENSOR_TENSOR_MODE, this->processDataNum);
        }
        
        inQueueX1.FreeTensor(x1Local);
        inQueueX2.FreeTensor(x2Local);
        outQueueY.EnQue<TYPE_Y>(yLocal);
    }
    __aicore__ inline void CopyOut(int32_t progress) {
        LocalTensor<DTYPE_Y> yLocal = outQueueY.DeQue<DTYPE_Y>();
        DataCopy(yGm[progress * this->tileDataNum], yLocal, this->processDataNum);
        outQueueY.FreeTensor(yLocal);
    }

private:
    TPipe* pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueCondition, inQueueX1, inQueueX2;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TBuf<QuePosition::VECCALC> QueueTmpbit, QueueTmp1, QueueTmp2;

    GlobalTensor<DTYPE_CONDITION> conditionGm;
    GlobalTensor<DTYPE_X1> x1Gm;
    GlobalTensor<DTYPE_X2> x2Gm;
    GlobalTensor<DTYPE_Y> yGm;

    uint16_t tileNum;
    uint16_t tileDataNum;
    uint16_t tailDataNum;
    uint16_t processDataNum;
};
template<typename TYPE_Y> class KernelSelectV2_broadcast {
    public:
        __aicore__ inline KernelSelectV2_broadcast () {}
        __aicore__ inline void Init(GM_ADDR condition, GM_ADDR x1, GM_ADDR x2, GM_ADDR y,
                                uint32_t size,uint32_t y_dim, uint32_t* indices_shape, uint32_t* x1_shape, uint32_t* x2_shape, uint32_t* y_shape) {
            ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
            this->size = size;
            this->y_dim = y_dim;
            for(uint32_t i=0; i<this->y_dim; i++)
            {
                this->indices_shape[i] = indices_shape[i];
                this->x1_shape[i] = x1_shape[i];
                this->x2_shape[i] = x2_shape[i];
                this->y_shape[i] = y_shape[i];
            }
            
            conditionGm.SetGlobalBuffer((__gm__ DTYPE_CONDITION*)condition, this->size);
            x1Gm.SetGlobalBuffer((__gm__ DTYPE_X1*)x1, this->size);
            x2Gm.SetGlobalBuffer((__gm__ DTYPE_X2*)x2, this->size);
            yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y, this->size);
    
        }
        __aicore__ inline void Process() {
            for(uint32_t i=0; i<this->size; i++)
            {
                if(conditionGm.GetValue(get_broadcasted_index(i, indices_shape, y_shape, y_dim)))
                {
                    yGm.SetValue(i, x1Gm.GetValue(get_broadcasted_index(i, x1_shape, y_shape, y_dim)));
                }
                else
                {
                    yGm.SetValue(i, x2Gm.GetValue(get_broadcasted_index(i, x2_shape, y_shape, y_dim)));
                }
            }
        }
    private:
    __aicore__ inline uint32_t get_broadcasted_index(uint32_t indices, uint32_t *original_shape, uint32_t *broadcast_shape, uint32_t ndim)
    {
        uint32_t shape[8];
        uint32_t original_indices = 0;
        int i =0;
        for(i=ndim-1; i>=0; i--)
        {
            shape[i] = 0;
            if(original_shape[i]!=1)
            {
                shape[i] = indices % broadcast_shape[i];
            }
            indices /= broadcast_shape[i];
            if(indices == 0) break;   
        }
        for(; i<ndim; i++)
        {
            original_indices = shape[i] + original_indices * original_shape[i];
        }
        return original_indices;
    }
    private:
    
        GlobalTensor<DTYPE_CONDITION> conditionGm;
        GlobalTensor<DTYPE_X1> x1Gm;
        GlobalTensor<DTYPE_X2> x2Gm;
        GlobalTensor<DTYPE_Y> yGm;
        uint32_t size;
    
        uint32_t y_dim;
        uint32_t indices_shape[8];
        uint32_t x1_shape[8];
        uint32_t x2_shape[8];
        uint32_t y_shape[8];
    };
extern "C" __global__ __aicore__ void select_v2(GM_ADDR condition, GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl

    if (TILING_KEY_IS(1)) {
        TPipe pipe;
        KernelSelectV2<DTYPE_X1> op;
        op.Init(condition, x1, x2, y, 
                tiling_data.finalTileNum, tiling_data.tileDataNum, tiling_data.TailDataNum, &pipe);
        op.Process();
    }
    else if (TILING_KEY_IS(2)) {
        KernelSelectV2_broadcast<DTYPE_X1> op;
        op.Init(condition, x1, x2, y, 
                tiling_data.size, tiling_data.y_dim, tiling_data.indices_shape, tiling_data.x1_shape, tiling_data.x2_shape, tiling_data.y_shape);
        op.Process();
    }
    
}